# 使用langchain框架去记录上下文---（Memory）

from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
from langchain.memory import ConversationBufferWindowMemory
from langchain.memory import ConversationTokenBufferMemory
from langchain.memory import ConversationSummaryBufferMemory
import warnings
warnings.filterwarnings('ignore')

api_key = "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"

# 对话缓冲内存记忆（ConversationBufferMemory）
def get_langchain_memory_1():
    chat = ChatOpenAI(api_key=api_key, temperature=0.0)
    memory = ConversationBufferMemory()
    conversation = ConversationChain(
        llm=chat,
        memory=memory,
        verbose=True
    )

    conversation.predict(input="Hi, my name is Andrew")
    conversation.predict(input="What is 1+1?")
    conversation.predict(input="What is my name?")
    print(memory.buffer)

    memory.save_context({"input": "Hi"},
                    {"output": "What's up"})
    memory.save_context({"input": "Not much, just hanging"},
                        {"output": "Cool"})

    intput = conversation.predict(input="Hi")
    print(intput)

# 窗口缓冲内存记忆（ConversationBufferWindowMemory）
def get_langchain_memory_2():
    chat = ChatOpenAI(api_key=api_key, temperature=0.0)
    # k=1代表只能记忆一个对话
    memory = ConversationBufferWindowMemory(k=1)
    conversation = ConversationChain(
        llm=chat,
        memory=memory,
        verbose=True
    )

    conversation.predict(input="Hi, my name is Andrew")
    conversation.predict(input="What is 1+1?")
    conversation.predict(input="What is my name?")
    print(memory.buffer)

    memory.save_context({"input": "Hi"},
                    {"output": "What's up"})
    memory.save_context({"input": "Not much, just hanging"},
                        {"output": "Cool"})

    intput = conversation.predict(input="Hi")
    print(intput)

# Token缓冲内存记忆（ConversationTokenBufferMemory）
def get_langchain_memory_3():
    chat = ChatOpenAI(api_key=api_key, temperature=0.0)
    # max_token_limit代表能记忆多少个Token（词向量）
    memory = ConversationTokenBufferMemory(llm=chat, max_token_limit=30)
    memory.save_context({"input": "AI is what?!"},
                        {"output": "Amazing!"})
    memory.save_context({"input": "Backpropagation is what?"},
                        {"output": "Beautiful!"})
    memory.save_context({"input": "Chatbots are what?"},
                        {"output": "Charming!"})
    print(memory.buffer)

# 摘要缓冲内存记忆（ConversationSummaryMemory）
def get_langchain_memory_4():
    schedule = "There is a meeting at 8am with your product team. \
    You will need your powerpoint presentation prepared. \
    9am-12pm have time to work on your LangChain \
    project which will go quickly because Langchain is such a powerful tool. \
    At Noon, lunch at the italian resturant with a customer who is driving \
    from over an hour away to meet you to understand the latest in AI. \
    Be sure to bring your laptop to show the latest LLM demo."
    chat = ChatOpenAI(api_key=api_key, temperature=0.0)
    # max_token_limit代表能记忆多少个摘要
    memory = ConversationSummaryBufferMemory(llm=chat, max_token_limit=200)
    conversation = ConversationChain(
        llm=chat,
        memory=memory,
        verbose=True
    )
    memory.save_context({"input": "Hello"}, {"output": "What's up"})
    memory.save_context({"input": "Not much, just hanging"},
                        {"output": "Cool"})
    memory.save_context({"input": "What is on the schedule today?"},
                        {"output": f"{schedule}"})

    conversation.predict(input="What would be a good demo to show?")
    print(memory.buffer)

if __name__ == '__main__':
    # get_langchain_memory_1()
    # get_langchain_memory_2()
    # get_langchain_memory_3()
    get_langchain_memory_4()